import numpy as np

# This class implements a writer for a restricted form of the Metis graph 
# format. This implementation assumes indices are base-0 in the graph, but
# writes them in base-1 (as per the metis spec).

# THIS ASSUMES THE ADJACENCY MATRIX of the graph IS SYMMETRIC.  IF NOT, METIS MAY FAIL.

class MetisGraphWriter:

    def __init__(self):
        self.stream = None
        self.graph = None

    def write(self, graph, filename):
        # now output the cmi values
        out = open(filename, 'w')
        n = len(graph)

        # divide the edge count by two because metis treats the graph as undirected
        edgeCount = graph.getEdgeCount() / 2

        if graph.weighted:
            fmt = "001"
        else:
            fmt = "000"
        out.write("{} {} {}\n".format(n, int(edgeCount), fmt))

        # the +1s are necessary because metis is base-1
        for x in range(n):
            for y in range(n):
                if graph.getEdge(x, y) != 0:
                    if graph.weighted:
                        out.write("{} {} ".format((y+1), int(graph.getEdge(x, y))))
                    else:
                        out.write("{} ".format((y+1)))
            out.write("\n")

        out.close()
